package com.oracle.coherence.patterns.security;

import com.tangosol.net.CacheFactory;
import com.tangosol.run.xml.XmlElement;

import javax.security.auth.Subject;
import javax.security.auth.SubjectDomainCombiner;
import javax.security.auth.kerberos.KerberosPrincipal;
import javax.security.auth.login.LoginContext;
import javax.security.auth.login.LoginException;
import java.security.AccessControlContext;
import java.security.AccessController;
import java.security.DomainCombiner;
import java.security.Principal;
import java.util.Set;
import java.util.TreeSet;

/**
 * Utility methods for obtaining a{@link javax.security.auth.Subject} that contains
 * a {@link KrbTicketPrincipal} with the correct
 * tickets to connect to a specific server service.
 * <p/>
 *
 * @author Jonathan Knight
 */
public class KrbSecurityHelper {
    public static final String SYSTEM_PROPERTY_SPN = "security.krb.servicename";

    /**
     * Obtains a {@link javax.security.auth.Subject} using the login module
     * specified in the Coherence operational configuration and containing an
     * instance of a {@link KrbTicketPrincipal} with the required tickets to connect and
     * authenticate to the service name specified using the JVM parameter
     * <code>-Dsecurity.krb.servicename</code>.
     * <p/>
     *
     * @return a {@link javax.security.auth.Subject} wrapping a
     *         {@link KrbTicketPrincipal} for connection to the
     *         specified service.
     * @throws LoginException - if an error occurs obtaining the security credentials.
     */
    public static Subject login() throws LoginException {
        String servicePrincipalName = System.getProperty(SYSTEM_PROPERTY_SPN);
        return login(servicePrincipalName);
    }

    /**
     * Obtains a {@link javax.security.auth.Subject} using the login module
     * specified in the Coherence operational configuration and containing an
     * instance of a {@link KrbTicketPrincipal} with the required tickets to connect and
     * authenticate to the specified service name.
     * <p/>
     *
     * @param servicePrincipalName - the name of the service to authenticate to.
     * @return a {@link javax.security.auth.Subject} wrapping a
     *         {@link KrbTicketPrincipal} for connection to the
     *         specified service.
     * @throws LoginException - if an error occurs obtaining the security credentials.
     */
    public static Subject login(String servicePrincipalName) throws LoginException {

        Subject subject = getCurrentSubject();
        if (subject == null) {
            XmlElement xmlCluster = CacheFactory.getSecurityConfig();
            XmlElement xmlLogin = xmlCluster.getSafeElement("login-module-name");
            LoginContext lc = new LoginContext(xmlLogin.getString("Coherence"));
            lc.login();
            subject = lc.getSubject();
        }

        return login(servicePrincipalName, subject);
    }

    /**
     * Obtains a {@link javax.security.auth.Subject} that wraps a
     * {@link KrbTicketPrincipal} containing the correct tickets
     * to connect to and authenticate to the specified service.
     * <p/>
     * The {@link javax.security.auth.Subject} specified in the subject parameter is used to obtain the
     * required tickets from Kerberos. These tickets will be wrapped in a
     * {@link KrbTicketPrincipal}.
     * <p/>
     * If the subject passed in is readOnly a new {@link javax.security.auth.Subject} will be created
     * containing the same credentials as the original plus the
     * {@link KrbTicketPrincipal} otherwise the
     * {@link KrbTicketPrincipal} will be added to the existing
     * {@link javax.security.auth.Subject}
     * <p/>
     *
     * @param servicePrincipalName - the name of the service to obtain tickets for
     * @param subject              - the {@link javax.security.auth.Subject} to use to obtain tickets
     *                             for the service.
     * @return A {@link javax.security.auth.Subject} containing the credentials required to authenticate
     *         to the specified service.
     * @throws LoginException - if an error occurs obtaining the security credentials.
     */
    public static Subject login(final String servicePrincipalName, Subject subject) throws LoginException {
        if (servicePrincipalName == null || servicePrincipalName.length() == 0) {
            throw new IllegalArgumentException("The servicePrincipalName must not be null or empty string");
        }

        if (subject == null) {
            throw new IllegalArgumentException("The Subject parameter cannot be null");
        }

        Set<KerberosPrincipal> krbPrincipals = subject.getPrincipals(KerberosPrincipal.class);
        final KerberosPrincipal krbPrincipal = krbPrincipals.iterator().next();

        final KrbTicketPrincipal principal = new KrbTicketPrincipal(krbPrincipal.getName(), servicePrincipalName);
        principal.setAuthenticated(true);
        principal.setLocal(true);

        if (subject.isReadOnly()) {
            Set<Principal> principals = new TreeSet<Principal>(subject.getPrincipals());
            principals.add(principal);
            subject = new Subject(true,
                    principals,
                    subject.getPublicCredentials(),
                    subject.getPrivateCredentials());
        } else {
            subject.getPrincipals().add(principal);
        }

        return subject;
    }

    /**
     * Return the Subject associated with the calling thread.
     *
     * @return the current Subject or null if a Subject is not associated with
     *         the calling thread
     */
    public static Subject getCurrentSubject() {
        Subject subject = null;
        Object context;

        SecurityManager manager = System.getSecurityManager();

        if (manager == null) {
            context = AccessController.getContext();
        } else {
            context = manager.getSecurityContext();
        }

        if (context != null && context instanceof AccessControlContext) {
            DomainCombiner dc = ((AccessControlContext) context).getDomainCombiner();
            if (dc != null && dc instanceof SubjectDomainCombiner) {
                subject = ((SubjectDomainCombiner) dc).getSubject();
            }
        }
        return subject;
    }

}
